# coding:utf-8

import csv
import os
import codecs
import time

import cv2
import numpy as np

from Utils import utils
from Utils.DataLoader import IITDLoader, CASIALoader, TongjiLoader, CASIANormalLoader
from Zhang.zhangROI import ZhangROIExtractor


def generate_empty_KP(image, file_name, writer):
    pts = dict()
    pts['file_name'] = file_name
    pts['x0'] = 0
    pts['x2'] = 480 - 1
    pts['y0'] = 0
    pts['y2'] = 640 - 1
    writer.writerow(pts)


def extractAndSaveKP_IITD(img_ori, image_crop_resize, file_name, path_save, rio_extractor, writer):
    img_save_path = '/home/yjy/dataset/IITD-KP-Zhang'
    img_gray = cv2.cvtColor(image_crop_resize, cv2.COLOR_RGB2GRAY)
    roi, pts = rio_extractor.extract(img_gray)
    pts['file_name'] = file_name
    # x towards down; y towards right
    x0 = float(pts['x0'])
    x2 = float(pts['x2'])
    y0 = float(pts['y0'])
    y2 = float(pts['y2'])
    # img_crop = img[52:428, 140:, :]
    # shape = np.shape(img_ori)
    # factor = (640 / shape[1], 480 / shape[0])
    # img_resize = cv2.resize(img_ori, factor)
    x0 = x0 / 1.28 + 52
    x2 = x2 / 1.28 + 52
    y0 = y0 / 1.28 + 140
    y2 = y2 / 1.28 + 140
    # pts['x0'] = y0
    # pts['x2'] = y2
    # pts['y0'] = -x0
    # pts['y2'] = -x2
    # cv2.imwrite(os.path.join(path_save, file_name), roi)
    if roi.sum() > 100:
        cv2.imwrite(os.path.join(path_save, file_name), roi)
        writer.writerow(pts)
        cv2.circle(img_ori, center=(int(y0), int(x0)), radius=5, color=(0, 0, 128), thickness=-1)
        cv2.circle(img_ori, center=(int(y2), int(x2)), radius=10, color=(0, 0, 128), thickness=-1)
        cv2.imwrite(os.path.join(img_save_path, file_name), img_ori)


def extractAndSaveKP_Tongji(image, file_name, path_save, rio_extractor, writer):
    img_save_path = '/home/yjy/dataset/Tongji-KP-NN'
    img_gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
    roi, pts = rio_extractor.extract(img_gray)
    pts['file_name'] = file_name
    # x towards down; y towards right
    x0 = float(pts['x0'])
    x2 = float(pts['x2'])
    y0 = float(pts['y0'])
    y2 = float(pts['y2'])
    if roi.sum() < 100:
        # if ROI wrong, use default coordinates
        roi = rio_extractor.extract_kp(img_gray, x0, x2, y0, y2)
    cv2.imwrite(os.path.join(path_save, file_name), roi)
    writer.writerow(pts)
    cv2.circle(image, center=(int(y0), int(x0)), radius=5, color=(0, 0, 128), thickness=-1)
    cv2.circle(image, center=(int(y2), int(x2)), radius=10, color=(0, 0, 128), thickness=-1)
    cv2.imwrite(os.path.join(img_save_path, file_name), image)


def extractAndSaveKP_CASIA(image, file_name, path_save, rio_extractor, writer):
    img_save_path = '/home/yjy/dataset/CASIA-KP-RotNN'
    # img_gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
    img_gray = image
    curr_time = time.time()
    roi, pts = rio_extractor.extract(img_gray)
    timecost = time.time() - curr_time
    pts['file_name'] = file_name
    # x towards down; y towards right
    x0 = float(pts['x0'])
    x2 = float(pts['x2'])
    y0 = float(pts['y0'])
    y2 = float(pts['y2'])

    if roi.sum() < 100:
        # if ROI wrong, use default coordinates
        roi = rio_extractor.extract_kp(img_gray, x0, x2, y0, y2)
    cv2.imwrite(os.path.join(path_save, file_name), roi)
    writer.writerow(pts)
    cv2.circle(image, center=(int(y0), int(x0)), radius=5, color=(0, 0, 128), thickness=-1)
    cv2.circle(image, center=(int(y2), int(x2)), radius=10, color=(0, 0, 128), thickness=-1)
    cv2.imwrite(os.path.join(img_save_path, file_name), image)
    return timecost


def extractZhangROI(dataLoader, rio_extractor):
    path_save = '/home/yjy/dataset/CASIA-ROI-RotNN'
    csv_path = '/home/yjy/dataset/casia_rot_kp_label.csv'
    with codecs.open(csv_path, 'w', 'utf-8') as csv_file:
        fieldnames = ['file_name', 'x0', 'x2', 'y0', 'y2']  # 表头
        writer = csv.DictWriter(csv_file, fieldnames=fieldnames)
        writer.writeheader()
        total_time = dataLoader.map(extractAndSaveKP_CASIA, path_save, rio_extractor, writer)
    print 'Mean ROI extraction cost time:', 1000 * total_time.mean(), ' ms'


def extractNNROI(img_paths, csv_path, roi_save_path, rio_extractor):
    img_save_path = '/home/yjy/dataset/CASIA-KP-NN'
    # dirs = os.listdir(img_paths)
    M_Time = []
    with codecs.open(csv_path, 'r', 'utf-8') as csv_file:
        reader = csv.DictReader(csv_file)
        for row in reader:
            # filename = row[0]
            # x0, x2, y0, y2 = map(lambda x: float(x), row[1:])
            filename = row['file_name']
            x0 = float(row['x0'])
            x2 = float(row['x2'])
            y0 = float(row['y0'])
            y2 = float(row['y2'])
            img = cv2.imread(os.path.join(img_paths, filename), cv2.IMREAD_GRAYSCALE)
            curr_time = time.time()
            roi = rio_extractor.extract_kp(img, x0, x2, y0, y2)
            M_Time.append(time.time() - curr_time)
            cv2.imwrite(os.path.join(roi_save_path, filename), roi)
            # center=(w,h)
            cv2.circle(img, center=(int(y0), int(x0)), radius=5, color=(0, 0, 128), thickness=-1)
            cv2.circle(img, center=(int(y2), int(x2)), radius=10, color=(0, 0, 128), thickness=-1)
            cv2.imwrite(os.path.join(img_save_path, filename), img)
    total_time = np.array(M_Time)
    print 'Mean ROI extraction cost time:', 1000 * total_time.mean(), ' ms'


def generate_all_csv(dataLoader, csv_path):
    with codecs.open(csv_path, 'w', 'utf-8') as csv_file:
        fieldnames = ['file_name', 'x0', 'x2', 'y0', 'y2']  # 表头
        writer = csv.DictWriter(csv_file, fieldnames=fieldnames)
        writer.writeheader()
        dataLoader.map(generate_empty_KP, writer)


if __name__ == '__main__':
    img_path = '/home/yjy/dataset/CASIA-rotation-NN'
    # img_path = '/home/yjy/dataset/CASIA_rotation_tradition'
    # img_path = '/home/yjy/dataset/palmprint_dectection_tongji/LHand/palmprint_rotation_adjusted/'
    # csv_path = '/home/yjy/dataset/casia_label.csv'
    csv_handcraft_path = '/home/yjy/PycharmProjects/HRNet-Facial-Landmark-Detection/data/casia/palm_landmarks_casia_all.csv'
    img_path_merge = '/home/yjy/PycharmProjects/HRNet-Facial-Landmark-Detection/data/casia/images'
    csv_NN_path = '/home/yjy/PycharmProjects/HRNet-Facial-Landmark-Detection/output/CASIA/palm_alignment_casia_hrnet_w18' \
                  '/predictions_all.csv'
    roi_save_path = '/home/yjy/dataset/CASIA-ROI-NN'

    roie = ZhangROIExtractor()
    dataLoader = CASIANormalLoader(img_path)
    # generate_all_csv(dataLoader, csv_handcraft_path)
    # extractZhangROI(dataLoader, roie)
    extractNNROI(img_path_merge, csv_NN_path, roi_save_path, roie)
